package reverse_proxy

import (
	"context"
	"io"
	"log"
	"net"
	"time"

	"go_gateway/reverse_proxy/load_balance"
	"go_gateway/tcp_proxy_middleware"
)

/*
TCP 反向代理的核心实现，负责将 TCP 连接从客户端代理到后端服务器
工作模式：
客户端 (src) ↔ 网关TCP代理 ↔ 后端服务器 (dst)

完整代理流程:
1. 客户端连接到网关TCP代理
2. 负载均衡选择后端服务器
3. 网关连接到后端服务器
4. 启动双向数据转发：
   ┌─ goroutine1: 客户端数据 → 后端服务器
   └─ goroutine2: 后端服务器数据 → 客户端
5. 任一方向连接断开，整个代理终止

*/

// 创建负载均衡代理
func NewTcpLoadBalanceReverseProxy(c *tcp_proxy_middleware.TcpSliceRouterContext, lb load_balance.LoadBalance) *TcpReverseProxy {
	return func() *TcpReverseProxy {
		nextAddr, err := lb.Get("") // 从负载均衡器获取后端地址
		if err != nil {
			log.Fatal("get next addr fail")
		}
		return &TcpReverseProxy{
			ctx:             c.Ctx,
			Addr:            nextAddr, // 设置后端服务器地址
			KeepAlivePeriod: time.Second,
			DialTimeout:     time.Second,
		}
	}()
}

// TCP反向代理
type TcpReverseProxy struct {
	ctx                  context.Context //单次请求单独设置
	Addr                 string          // 后端服务器地址
	KeepAlivePeriod      time.Duration   // 保活周期
	DialTimeout          time.Duration   //设置超时时间
	DialContext          func(ctx context.Context, network, address string) (net.Conn, error)
	OnDialError          func(src net.Conn, dstDialErr error) // 拨号错误回调
	ProxyProtocolVersion int                                  // 代理协议版本
}

// 超时控制
func (dp *TcpReverseProxy) dialTimeout() time.Duration {
	if dp.DialTimeout > 0 {
		return dp.DialTimeout
	}
	return 10 * time.Second // 默认10秒
}

var defaultDialer = new(net.Dialer)

// 拨号器配置
func (dp *TcpReverseProxy) dialContext() func(ctx context.Context, network, address string) (net.Conn, error) {
	if dp.DialContext != nil {
		return dp.DialContext
	}
	return (&net.Dialer{
		Timeout:   dp.DialTimeout,     //连接超时
		KeepAlive: dp.KeepAlivePeriod, //设置连接的检测时长
	}).DialContext
}

// 保活设置
func (dp *TcpReverseProxy) keepAlivePeriod() time.Duration {
	if dp.KeepAlivePeriod != 0 {
		return dp.KeepAlivePeriod
	}
	return time.Minute // 默认1分钟
}

// 处理 TCP 代理: 传入上游 conn，在这里完成下游连接与数据交换
func (dp *TcpReverseProxy) ServeTCP(ctx context.Context, src net.Conn) {
	// 设置连接后端的超时
	var cancel context.CancelFunc
	if dp.DialTimeout >= 0 {
		ctx, cancel = context.WithTimeout(ctx, dp.dialTimeout())
	}
	// 1.建立到后端服务器的连接（第二次握手：代理服务器 ↔ 后端服务器）
	dst, err := dp.dialContext()(ctx, "tcp", dp.Addr) // 连接后端服务器
	if cancel != nil {
		cancel()
	}
	if err != nil {
		dp.onDialError()(src, err)
		return
	}

	// 2.确保后端连接关闭
	defer func() { go dst.Close() }() //记得退出下游连接

	// 3.设置后端连接保活: 设置dst的 keepAlive 参数,在数据请求之前
	if ka := dp.keepAlivePeriod(); ka > 0 {
		if c, ok := dst.(*net.TCPConn); ok {
			c.SetKeepAlive(true)
			c.SetKeepAlivePeriod(ka)
		}
	}

	// 4.双向数据转发
	errc := make(chan error, 1)
	go dp.proxyCopy(errc, src, dst) // 后端 → 客户端：从dst(后端)读，往src(客户端)写
	go dp.proxyCopy(errc, dst, src) // 客户端 → 后端：从src(客户端)读，往dst(后端)写
	<-errc                          // 等待任意一个方向出错
}

// 连接失败处理
func (dp *TcpReverseProxy) onDialError() func(src net.Conn, dstDialErr error) {
	if dp.OnDialError != nil {
		return dp.OnDialError
	}
	return func(src net.Conn, dstDialErr error) {
		log.Printf("tcpproxy: for incoming conn %v, error dialing %q: %v", src.RemoteAddr().String(), dp.Addr, dstDialErr)
		src.Close()
	}
}

// 双向数据复制
func (dp *TcpReverseProxy) proxyCopy(errc chan<- error, dst, src net.Conn) {
	_, err := io.Copy(dst, src) // 将 src 的数据复制到 dst
	errc <- err                 // 发送复制结果
}
